package com.czw.stream_ai.web;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

@Controller
@CrossOrigin(origins = "*")
public class StreamController {
    @Value("${apikey}")
    private String apiKey;

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter streamGet(@RequestParam("userInput") String userInput) {
        return streamCommon(userInput);
    }

    private SseEmitter streamCommon(String userInput) {
        SseEmitter emitter = new SseEmitter();

        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        ArkService service =
                ArkService.builder()
                        .dispatcher(dispatcher)
                        .connectionPool(connectionPool)
                        .apiKey(apiKey)
                        .build();

        List<ChatMessage> messagesForReqList = new ArrayList<>();
        final ChatMessage systemMessage = ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content("你是程志炜开发的AI人工智能助手，你的名字是志炜一言。请遵循以下回复规则：\n" +
                        "1. 英文单词之间必须添加空格\n" +
                        "2. 中文内容使用Markdown格式进行排版\n" +
                        "3. 长文本合理分段，使用###作为小标题\n" +
                        "4. 保持回复专业且友好")
                .build();

        final ChatMessage userMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(userInput)
                .build();

        messagesForReqList.add(systemMessage);
        messagesForReqList.add(userMessage);

        ChatCompletionRequest req =
                ChatCompletionRequest.builder()
                        .model("doubao-1-5-pro-32k-250115")
                        .messages(messagesForReqList)
                        .stream(true)
                        .build();

        service.streamChatCompletion(req)
                .doOnError(throwable -> {
                    emitter.completeWithError(throwable);
                    service.shutdownExecutor();
                })
                .subscribe(choice -> {
                    if (choice.getChoices().size() > 0) {
                        try {
                            String content = (String) choice.getChoices().get(0).getMessage().getContent();
                            content = formatContent(content);
                            emitter.send(SseEmitter.event().data(content));
                        } catch (Exception e) {
                            emitter.completeWithError(e);
                            service.shutdownExecutor();
                        }
                    }
                }, throwable -> {
                    emitter.completeWithError(throwable);
                    service.shutdownExecutor();
                }, () -> {
                    emitter.complete();
                    service.shutdownExecutor();
                });

        return emitter;
    }

    private String formatContent(String content) {
        // 1. 处理英文空格
        content = content.replaceAll("([a-z])([A-Z])", "$1 $2")
                .replaceAll("([A-Z])([A-Z][a-z])", "$1 $2");

        // 2. 处理中文Markdown格式
        content = content.replace("###", "\n\n### ")
                .replace("\n", "\n\n");

        // 3. 确保标点符号规范
        content = content.replace("，", ", ")
                .replace("。", ". ")
                .replaceAll("\\s+", " ")
                .trim();

        return content;
    }
}